package com.chengqs.leetcode.hot100;

import com.chengqs.leetcode.utils.DataGeneratorUtil;
import com.chengqs.leetcode.utils.TimeCostUtil;

import java.util.HashMap;
import java.util.Map;

/**
 * 子串<br>
 * 数组、哈希表、前缀和<br>
 * 中等
 *
 * <h1>560. 和为 K 的子数组</h1>
 *
 * <p>给你一个整数数组 nums 和一个整数 k ，请你统计并返回 该数组中和为 k 的子数组的个数 。</p>
 *
 * <p>子数组是数组中元素的连续非空序列。</p>
 */
public class D10SubarraySum {
    public static void main(String[] args) {
        D10SubarraySum subarraySum = new D10SubarraySum();
        int[] nums = {1, 2, 3};

        TimeCostUtil.timeCost("枚举", () -> subarraySum.solution1(nums, 3));
        TimeCostUtil.timeCost("前缀和 + 哈希表优化", () -> subarraySum.solution2(nums, 3));
        TimeCostUtil.timeCost("前缀和 + 哈希表优化_我的的方法", () -> subarraySum.solution2_mine(nums, 3));

        System.out.println("大数据量");
        int[] nums2 = DataGeneratorUtil.generateRandomIntArray(100000, -50, 50);
        TimeCostUtil.timeCost("前缀和 + 哈希表优化", () -> subarraySum.solution2(nums2, 50));
        TimeCostUtil.timeCost("前缀和 + 哈希表优化_我的的方法", () -> subarraySum.solution2_mine(nums2, 50));
        TimeCostUtil.timeCost("枚举", () -> subarraySum.solution1(nums2, 50));

    }

    public int solution2_mine(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>(); // 统计每个前缀和出现的次数
        // 前缀和
        int sum = 0;
        // 统计
        int count = 0;
        map.put(0, 1);
        for (int i = 0; i < nums.length; i++) {
            sum += nums[i];
            if (map.containsKey(sum - k)) {
                count += map.get(sum - k); // 统计计数
            }
            map.put(sum, map.getOrDefault(sum, 0) + 1); // 如果从来没出现过，则默认为0

        }
        return count;
    }

    // 枚举
    public int solution1(int[] nums, int k) {
        int count = 0;
        for (int start = 0; start < nums.length; ++start) {
            int sum = 0;
            for (int end = start; end >= 0; --end) {
                sum += nums[end];
                if (sum == k) {
                    count++;
                }
            }
        }
        return count;
    }

    // 前缀和 + 哈希表优化
    public int solution2(int[] nums, int k) {
        int count = 0, pre = 0;
        HashMap<Integer, Integer> mp = new HashMap<>();
        mp.put(0, 1);
        for (int i = 0; i < nums.length; i++) {
            pre += nums[i];
            if (mp.containsKey(pre - k)) {
                count += mp.get(pre - k);
            }
            mp.put(pre, mp.getOrDefault(pre, 0) + 1);
        }
        return count;
    }
}
